
from matplotlib.pyplot import figure
import pandas as pd
import plotly.graph_objects as go
import plotly.express as px
from plotly.graph_objs.layout import xaxis


# tjd_legend_opts = opts.LegendOpts(pos_top="7%", is_show=True, textstyle_opts=opts.TextStyleOpts(font_size=18))
# tjd_text_opts = opts.TextStyleOpts(font_size=18)


def multi_macro_bar_or_line(df, chart_type, title_name, add_text=False):
    """
    绘制多条线的柱形图或线图
    """
    if chart_type == 'bar':
        # 对于柱状图，添加数据标签
        # fig = px.bar(df, x=df.index, y=df.columns, barmode='group', text_auto=True)
        if add_text:
            fig = go.Figure(
            data=[go.Bar(x=df.index, y=df[col], name=col, text=df[col], textposition='auto') for col in df.columns],
            )
        else:
            fig = go.Figure(
                data=[go.Bar(x=df.index, y=df[col], name=col) for col in df.columns],
            )
    else:
        # 对于线图，通过 update_traces 方法添加数据标签
        if add_text:
            fig = go.Figure(
                data=[go.Scatter(x=df.index, y=df[col], name=col, mode='lines+markers+text', text=df[col], textposition='top center') for col in df.columns],
            )
        else:
            fig = go.Figure(
                data=[go.Scatter(x=df.index, y=df[col], name=col, mode='lines+markers', 
                                marker=dict(
                                    size=10,
                                    symbol='circle',
                                    )) for col in df.columns],
            )

    fig.update_layout(
        title=title_name,
        legend=dict(
            orientation="h",  # 水平排列
            yanchor="bottom",  # 垂直锚点为底部
            y=1.02,  # 靠近图表顶部
            xanchor="center",  # 水平锚点为中心
            x=0.5),  # 水平居中),
        xaxis_title="日期",
        yaxis_title="%",
        )
    return fig


def multi_grid_line(df):
    """
    绘制四方格线型图
    """
    df = df.dropna(how='all')
    df = df.round(2)

    fig_all = []
    for i, y_name in enumerate(df.columns.to_list()):
        se = df[y_name]
        se = se.dropna()
        se = se[-100:]
        se_df = pd.DataFrame(se)
        
        fig = px.line(se_df, x=se_df.index, y=y_name)
        fig.update_layout(legend=dict(orientation="h", yanchor="top", y=-0.1, xanchor="center", x=0.5))
        fig.update_traces(textposition="top center", text=df.values.tolist())
        fig_all.append(fig)
    
    return fig_all


def year_timeline_seasonal(df):
    # 轮播图，按年播放
    cols = df.columns.to_list()
    
    df = df.dropna(how='all')
    df = df.round(2)
    df.index = pd.to_datetime(df.index)

    years = df.index.year.unique().to_list()
    year_dic = {y_name:to_monthly_mean(df[y_name]) for y_name in cols}
    
    tml = Timeline(init_opts=opts.InitOpts(width='1400px', height='900px'))
    for year in years:
        b = Line()
        b.add_xaxis([f'{i}月' for i in range(1, 13)])
        for y_name in cols:
            seasonal_se = year_dic[y_name][year]
            b.add_yaxis(y_name, seasonal_se.to_list())
            b.set_global_opts(title_opts=opts.TitleOpts(f"{year}年"),
                              legend_opts=tjd_legend_opts
                              )
        tml.add(b, f"{year}年")
    return tml


def item_timeline_seasonal(df, freq, year_dic):
    # 轮播图，按数据播放（展示季节性, 数据为日度或周度或月）

    cols = df.columns.to_list()
    
    df = df.dropna(how='all')
    df = df.round(2)
    df.index = pd.to_datetime(df.index)

    years = df.index.year.unique().to_list()
    
    tml = Timeline(init_opts=opts.InitOpts(width='1400px', height='600px'))
    for y_name in cols:
        b = Line()
        
        if freq == 'D':
            xaxis_series = year_dic[y_name].index.to_list()
            smooth = True
        elif freq == 'M':
            xaxis_series = [f'{i}月' for i in range(1, 13)]
            smooth = False
        b.add_xaxis(xaxis_series)
        
        for year in years:
            seasonal_se = year_dic[y_name][year]
            if isinstance(seasonal_se, pd.DataFrame):
                seasonal_se = seasonal_se.iloc[:,0]
            if len(seasonal_se.dropna()) >= 1:
                if year == 2024:
                    b.add_yaxis(
                        f"{year}年", seasonal_se.to_list(),
                        linestyle_opts=opts.LineStyleOpts(width=2, color='black'),
                        itemstyle_opts=opts.ItemStyleOpts(color='black'),  # 设置数据点颜色
                        is_connect_nones=True, is_smooth=smooth,
                        label_opts=(opts.LabelOpts(is_show=False))
                        )
                else:
                    b.add_yaxis(f"{year}年", seasonal_se.to_list(),
                                is_connect_nones=True, is_smooth=smooth,
                                label_opts=(opts.LabelOpts(is_show=False)))
                
                b.set_global_opts(
                    title_opts=opts.TitleOpts(y_name),
                    legend_opts=opts.LegendOpts(pos_right="3%", pos_top="12%",
                                                orient='vertical',
                                                textstyle_opts=tjd_text_opts),
                    toolbox_opts=opts.ToolboxOpts(),
                                  )
            
        tml.add(b, y_name)
    return tml


def xy_match_same_freq_plot(df_cot, idx_fill, tztd_ratio):
    cols = df_cot.columns.to_list()
    b = Line()
    idx = [str(s)[:10] for s in df_cot.index]
    b.add_xaxis(idx)
    b.add_yaxis(cols[0], df_cot[cols[0]].to_list(), yaxis_index=0,
                color='black') # 添加第一条曲线，默认使用左侧y轴
    # 扩展一个右侧的y轴
    b.extend_axis(
        yaxis=opts.AxisOpts(name="特征X", type_="value", position="right")
    )
    b.add_yaxis(cols[1], df_cot[cols[1]], yaxis_index=1, color='red') # 添加第二条曲线，指定使用右侧y轴
    b.set_series_opts(
        markarea_opts=opts.MarkAreaOpts(
            data=[opts.MarkAreaItem(x=s) for s in idx_fill],
            itemstyle_opts=opts.ItemStyleOpts(color='pink',opacity=0.1),
        ),
                      )
    # 全局配置项
    b.set_global_opts(
        datazoom_opts=[
            opts.DataZoomOpts(range_start=10, range_end=100),
            opts.DataZoomOpts(type_="inside")],
        title_opts=opts.TitleOpts(title=f"因子关联分析>同涨同跌率:{100*tztd_ratio:.2f}%"),
        yaxis_opts=opts.AxisOpts(name="指标Y"),
        toolbox_opts=opts.ToolboxOpts(),
        tooltip_opts=opts.TooltipOpts(trigger="axis", axis_pointer_type="cross"),
        legend_opts=opts.LegendOpts(pos_top="6%", is_show=True)
    )
    return b